package sm2

import (
	"crypto/ecdsa"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"encoding/pem"
	"fmt"

	"codeup.aliyun.com/lsxd/zltx/adapter/utils/formater"
	"github.com/emmansun/gmsm/sm2"
	"github.com/emmansun/gmsm/smx509"
)

type KeyType int

const (
	Raw KeyType = iota + 1
	Hex
	Base64
	Pem
)

type Options struct {
	format     formater.Formater
	keyType    KeyType
	privateKey string
	publicKey  string
	enOpts     *sm2.EncrypterOpts
	deOpts     *sm2.DecrypterOpts
}

// SetFormat 设置数据格式
func (o *Options) SetFormat(f formater.Formater) {
	o.format = f
}

// SetKeyType 设置密钥格式
func (o *Options) SetKeyType(t KeyType) {
	o.keyType = t
}

// SetPrivateKey 设置私钥
func (o *Options) SetPrivateKey(key string) {
	o.privateKey = key
}

// SetPublicKey 设置公钥
func (o *Options) SetPublicKey(key string) {
	o.publicKey = key
}

// SetEncryptOpts 设置加密选项
func (o *Options) SetEncryptOpts(opts *sm2.EncrypterOpts) {
	o.enOpts = opts
}

// SetDecryptOpts 设置解密选项
func (o *Options) SetDecryptOpts(opts *sm2.DecrypterOpts) {
	o.deOpts = opts
}

// merge 合并选项
func (o *Options) merge(n *Options) {
	// 空值不合并
	if n == nil {
		return
	}
	if n.format != nil {
		o.format = n.format
	}
	if n.keyType > 0 {
		o.keyType = n.keyType
	}
	if n.privateKey != "" {
		o.privateKey = n.privateKey
	}
	if n.publicKey != "" {
		o.publicKey = n.publicKey
	}
	if n.enOpts != nil {
		o.enOpts = n.enOpts
	}
	if n.deOpts != nil {
		o.deOpts = n.deOpts
	}
}

type SM2 struct {
	o *Options
}

// NewSm2 创建一个新的SM2加密器
func NewSm2(o *Options) *SM2 {
	opt := &Options{
		format:  formater.Base64,
		keyType: Hex,
	}
	opt.merge(o)
	return &SM2{o: opt}
}

// Encrypt 加密数据
func (d *SM2) Encrypt(data string) (string, error) {
	if len(data) == 0 {
		return "", nil
	}
	content := []byte(data)

	publicKey, err := ParsePublicKey(d.o.publicKey, d.o.keyType)
	if err != nil {
		return "", err
	}

	dst, err := sm2.Encrypt(rand.Reader, publicKey, content, d.o.enOpts)
	if err != nil {
		return "", err
	}

	return d.o.format.Encode(dst), nil
}

// Decrypt 解密数据
func (d *SM2) Decrypt(data string) (string, error) {
	content, err := d.o.format.Decode(data)
	if err != nil {
		return "", err
	}
	if len(content) == 0 {
		return "", nil
	}
	privateKey, err := ParsePrivateKey(d.o.privateKey, d.o.keyType)
	if err != nil {
		return "", err
	}
	dst, err := privateKey.Decrypt(rand.Reader, content, d.o.deOpts)
	return string(dst), err
}

const publicPemFormat = `-----BEGIN PUBLIC KEY-----
%s
-----END PUBLIC KEY-----`

const privatePemFormat = `-----BEGIN PRIVATE KEY-----
%s
-----END PRIVATE KEY-----`

func ParsePublicKey(key string, kt KeyType) (*ecdsa.PublicKey, error) {
	if len(key) == 0 {
		return nil, nil
	}

	switch kt {
	case Raw:
		return sm2.NewPublicKey([]byte(key))
	case Hex:
		pubBytes, err := hex.DecodeString(key)
		if err != nil {
			return nil, err
		}
		return sm2.NewPublicKey(pubBytes)
	case Base64:
		pubBytes, err := base64.StdEncoding.DecodeString(key)
		if err != nil {
			return nil, err
		}
		return sm2.NewPublicKey(pubBytes)
	case Pem:
		pubPEM := fmt.Sprintf(publicPemFormat, key)
		block, _ := pem.Decode([]byte(pubPEM))
		if block == nil {
			return nil, fmt.Errorf("failed to parse PEM block containing the public key")
		}

		pub, err := smx509.ParsePKIXPublicKey(block.Bytes)
		if err != nil {
			return nil, err
		}

		if !sm2.IsSM2PublicKey(pub) {
			return nil, fmt.Errorf("not a SM2 public key")
		}

		if key, ok := pub.(*ecdsa.PublicKey); ok {
			return key, nil
		}

		return nil, fmt.Errorf("not a valid SM2 public key")
	default:
		return nil, nil
	}
}

func ParsePrivateKey(key string, kt KeyType) (*sm2.PrivateKey, error) {
	if len(key) == 0 {
		return nil, nil
	}

	switch kt {
	case Raw:
		return sm2.NewPrivateKey([]byte(key))
	case Hex:
		priBytes, err := hex.DecodeString(key)
		if err != nil {
			return nil, err
		}
		return sm2.NewPrivateKey(priBytes)
	case Base64:
		priBytes, err := base64.StdEncoding.DecodeString(key)
		if err != nil {
			return nil, err
		}
		return sm2.NewPrivateKey(priBytes)
	case Pem:
		priPEM := fmt.Sprintf(privatePemFormat, key)
		block, _ := pem.Decode([]byte(priPEM))
		if block == nil {
			return nil, fmt.Errorf("failed to parse PEM block containing the private key")
		}

		pri, err := smx509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, err
		}

		if pk, ok := pri.(*sm2.PrivateKey); ok {
			return pk, nil
		}
		return nil, fmt.Errorf("not a valid SM2 private key")
	default:
		return nil, fmt.Errorf("unsupported key type")
	}
}
